package com.example.demo.service;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Map;

@Slf4j
@Service
public class OllamaService {

    private final WebClient webClient;

    public OllamaService(@Value("${ollama.api.url:http://localhost:11434}") String ollamaApiUrl) {
        this.webClient = WebClient.builder()
                .baseUrl(ollamaApiUrl)
                .build();
    }

    /**
     * 生成文本
     */
    public Mono<String> generateText(String model, String prompt, Double temperature) {
        Map<String, Object> requestBody = Map.of(
                "model", model,
                "prompt", prompt,
                "stream", false
        );

        return webClient.post()
                .uri("/api/generate")
                .contentType(MediaType.APPLICATION_JSON)
                .body(BodyInserters.fromValue(requestBody))
                .retrieve()
                .bodyToMono(String.class) // 直接接收为字符串
                .map(response -> {
                    // 如果API实际返回JSON，可以在这里手动解析
                    // 但根据错误信息，可能不需要解析
                    return response;
                });
    }

    /**
     * 聊天对话
     */
    public Mono<String> chat(String model, List<Map<String, String>> messages) {
        Map<String, Object> requestBody = Map.of(
                "model", model,
                "messages", messages,
                "stream", false
        );

        return webClient.post()
                .uri("/api/chat")
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(requestBody)
                .retrieve()
                .bodyToMono(String.class) // 先接收为字符串
                .map(response -> {
//                    return response;   返回
                    // 如果API实际返回JSON，可以在这里手动解析
                    try {
                        ObjectMapper mapper = new ObjectMapper();
                        JsonNode root = mapper.readTree(response);

                        // 根据实际格式调整路径
                        if (root.has("message") && root.get("message").has("content")) {
                            return root.get("message").get("content").asText();
                        } else if (root.has("content")) {
                            return root.get("content").asText();
                        } else {
                            // 处理其他可能的格式
                            return response; // 如果无法解析，返回原始内容
                        }
                    } catch (Exception e) {
                        // 解析失败，可能不是JSON格式，直接返回原始内容
                        System.err.println("解析响应失败，将返回原始内容: " + e.getMessage());
                        return response;
                    }
                });
    }

    /**
     * 流式生成文本
     */
    public Flux<String> streamGenerateText(String model, String prompt, Double temperature) {
        Map<String, Object> requestBody = Map.of(
                "model", model,
                "prompt", prompt,
                "stream", true
        );

        return webClient.post()
                .uri("/api/generate")
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(requestBody)
                .retrieve()
                .bodyToFlux(String.class) // 接收为字符串流
                .filter(line -> !line.isEmpty()) // 过滤空行
                .map(line -> {
//                    return  line;
                    try {
                        // 手动解析每行JSON
                        ObjectMapper mapper = new ObjectMapper();
                        JsonNode node = mapper.readTree(line);

                        // 根据Ollama实际响应格式调整
                        if (node.has("response")) {
                            return node.get("response").asText();
                        } else if (node.has("done") && node.get("done").asBoolean()) {
                            return ""; // 流结束标记，返回空字符串或特殊标记
                        }
                        return ""; // 未知格式，返回空字符串
                    } catch (Exception e) {
                        System.err.println("解析行失败: " + line + ", 错误: " + e.getMessage());
                        return ""; // 解析失败，返回空字符串
                    }
                })
                .filter(chunk -> !chunk.isEmpty()); // 过滤解析后的空字符串
    }
}
